package multi_mutex

import frameworkRecover "payme3000/framework/recover"

// MultiMutex 通过channel模拟批量mutex
// example:
//
//	...
//	uid := 10086
//	m := NewMultiMutex()
//	func() {
//	    defer m.Lock(uid)()
//	    ...
//	}()
//	...
type MultiMutex struct {
	size int64 // mutex数量
	ch   []chan struct{}
}

func NewMultiMutex(opts ...int64) *MultiMutex {
	size := int64(0)

	// opts实际上只需要一个可选参数 size, 决定mutex数量
	if len(opts) > 0 {
		size = opts[0]
	}

	if size < 32 {
		size = 32
	}

	m := &MultiMutex{
		size: size,
	}

	// 生成channel, 每个缓冲长度1, 且初始化后让channel有数据, 处于可Lock状态
	for ; size > 0; size-- {
		ch := make(chan struct{}, 1)
		ch <- struct{}{}
		m.ch = append(m.ch, ch)
	}

	return m
}

// Lock 加锁
// @return unlock 调用后执行对应key的unlock
func (m *MultiMutex) Lock(key int64) (unLock func()) {
	// 计算出要使用哪个channel
	bit := key % m.size
	// 等待channel有数据
	<-m.ch[bit]
	// lock成功后记录版本号
	_ver := bit

	return func() {
		// 防止重复调用unlock, 版本号相同才执行unlock操作
		if _ver == bit {
			frameworkRecover.WithRecover(func() {
				defer func() {
					_ver++
				}()

				m.ch[bit] <- struct{}{}
			})
		}
	}
}
